# train.py
import subprocess

from fastapi import APIRouter, Form
from starlette.responses import HTMLResponse, JSONResponse
from pathlib import Path as SysPath

router = APIRouter()
@router.post("/train/")
async def train_concept_api(concept: str = Form(...)):
    images_dir = SysPath("./uploads") / concept
    output_dir = SysPath("./models") / concept
    output_dir.mkdir(parents=True, exist_ok=True)

    if not images_dir.exists():
        return JSONResponse({"error": f"未找到概念数据目录: {images_dir}"}, status_code=404)

    try:
        subprocess.run([
            "python", "train_lora.py",
            "--instance_data_dir", str(images_dir),
            "--output_dir", str(output_dir),
            "--concept_name", concept
        ], check=True)
        return {"status": "训练完成", "concept": concept}
    except subprocess.CalledProcessError as e:
        return JSONResponse({"error": f"训练失败: {e}"}, status_code=500)
